################################################################################

class rap_config:
  """ Stores a single configuration. """
  def __init__(self,conforne,ne=None):
    if isinstance(conforne,rap_jconfig):
      tmp=conforne.to_config()
      self.v_conf=tmp.v_conf
    elif isinstance(conforne,int):
      self.v_conf=self._baseconfig(conforne)
    else:
      self.v_conf=self._format_config(conforne,ne=ne)

# ------------------------------------------------------------------------------

  def __repr__(self):
    return self.label()

# ------------------------------------------------------------------------------

  def __eq__(self,conf):
    tconf=conf
    if not isinstance(tconf,rap_config): tconf=rap_config(tconf)
    if repr(self) == repr(tconf): return True
    return False

# ------------------------------------------------------------------------------

  def __ne__(self,conf):
    tconf=conf
    if not isinstance(tconf,rap_config): tconf=rap_config(tconf)
    if repr(self) != repr(tconf): return True
    return False

# ------------------------------------------------------------------------------

  def __cmp__(self,conf):
    return rlf_confsorter(self,conf)

# ------------------------------------------------------------------------------

  def __hash__(self):
    return hash(self.label())

# ------------------------------------------------------------------------------

  def __len__(self):
    """ Return length of configuration label. """
    return len(self.label())

# ------------------------------------------------------------------------------

  def __add__(self,other):
    return self.label()+other

# ------------------------------------------------------------------------------

  def __radd__(self,other):
    return other+self.label()

# ------------------------------------------------------------------------------

  def __getitem__(self,idx):
    return self.occupancy(idx)

# ------------------------------------------------------------------------------

  def __iter__(self):
    return self.subshells().__iter__()

# ------------------------------------------------------------------------------

  def _baseconfig(self,ne):
    """ Fill sub-shells in order until number of electrons is reached. """
    netot=ne
    out={}
    i=1
    while netot > 0:
      w=rlf_ssoccupancy(i)
      out[i]=min(w,netot)
      netot-=w
      i+=1
    return out

# ------------------------------------------------------------------------------

  def _format_config(self,inconf,ne=None):
    """ Attempt to understand input configuration. """

    # function to see if configuration is in Eissner form
    def read_eissner(x):
      x=str(x)
      tmp=re.findall(r'(\d+)([a-z])',x)
      if tmp: return False
      if len(x)%3 == 1: return False
      if len(x)%3 == 2: x='5'+x
      out={}
      for i in range(len(x)/3):
        xx=x[3*i:3*i+3]
        try:
          w=int(xx[:2])-50
          ss=rlf_eissner2ss(xx[-1],index=False)
        except:
          return False
        if w > 1 and rlf_ssindex(ss) > 9: return False   # probably not Eissner
        if w < 0 or w > rlf_ssoccupancy(ss): return False
        out[ss]=w
      conf=[]
      for ss,w in out.iteritems(): conf.append(ss+str(w))
      conf=' '.join(conf)
      return conf


    # configuration must be a string
    if isinstance(inconf,rap_config):
      return inconf
    if not isinstance(inconf,str): 
      msg='configuration must be a string; conf: %s'%inconf
      raise ErrorRAP(msg)

    conf=inconf+' '
    out={}

    # look for Eissner notation
    tmp=re.split(r'(\w+\.)',conf.strip())
    while tmp[0] == '': tmp.pop(0)
    chk=read_eissner(tmp[0])
    if chk is not False: conf=chk

    # configuration must contain both numbers and letters
    tmp=re.findall(r'(\d)',conf)
    if tmp == []: 
      msg='configuration contains no digits; conf: %s'%conf
      raise ErrorRAP(msg)
    tmp=re.findall(r'([A-Za-z])',conf)
    if tmp == []: 
      msg='configuration contains no letters; conf: %s'%conf
      raise ErrorRAP(msg)

    # if only uppercase letters exist, assume no term labels
    tmp=re.findall(r'([a-z])',conf)
    if tmp == []: conf=conf.lower()

    # get rid of carets
    conf=conf.replace('^','')

    # remove parity characters
    conf=conf.replace('e','')
    conf=conf.replace('o','')

    # identify any term labels and remove
    tmp=re.findall(r'(#\d+)',conf)
    for t in tmp: conf=conf.replace(t,'')
    tmp=re.findall(r'(\(?\s*\d+[A-Z]_?\d+\.5[^a-z]\s*\)?)',conf)
    for t in tmp: conf=conf.replace(t,'')
    tmp=re.findall(r'(\(?\s*\d+[A-Z]_?[\d/]*\s*\)?)',conf)
    for t in tmp: conf=conf.replace(t,'')
    conf=conf.strip()+' '

    # deal with hole-notation
    sshole=None
    whole=None
    idxhole=None
    tmp=re.findall(r'(\[\d+[a-z][\d]*\])',conf)
    if tmp != []:
      if len(tmp) > 1:
        msg='only one hole supported; conf: %s'%conf
        raise ErrorRAP(msg)
      t=tmp[0]
      ttmp=re.findall(r'(\d+)([a-z])(\d*)',t)
      n=int(ttmp[0][0])
      l=SPECLETTERS.index(ttmp[0][1].upper())
      w=1
      if ttmp[0][2].strip() != '': w=int(ttmp[0][2])
      wmax=4*l+2
      sshole=str(n)+SPECLETTERS[l].lower()
      whole=wmax-w
      if whole < 0: 
        msg='too many electrons in hole; conf: %s'%conf
        raise ErrorRAP(msg)
      conf=conf.replace(t,'')
      idxhole=rlf_ssindex(sshole)

    # check for letters and numbers
    tmp=re.findall(r'(\d)',conf)
    if len(tmp) == 0:
      msg='missing numbers in main part of configuration: %s'%inconf
      raise ErrorRAP(msg)
    tmp=re.findall(r'([A-Za-z])',conf)
    if len(tmp) == 0:
      msg='missing letters in main part of configuration: %s'%inconf
      raise ErrorRAP(msg)

    # need to separate sub-shells; first check for division characters
    # (space, period, etc) and if unsuccessful, do it the hard way
    tmp=re.split(r'[\W\s._]+',conf)
    while tmp[-1] == '': tmp.pop()
    while tmp[0] == '': tmp.pop(0)
    hardway=False
    ttmp=re.findall(r'([a-z])',conf)
    if len(ttmp) != len(tmp): hardway=True
    if not hardway:
      sslist=[]
      for t in tmp:
        if t == '': continue
        ttmp=list(re.findall(r'(\d+)([a-z])(\d*)',t)[0])
        if ttmp == []: 
          msg='cannot resolve sub-shells; conf: %s'%conf
          raise ErrorRAP(msg)
        if ttmp[-1] == '': ttmp[-1]=1
        sslist.append(ttmp)      
    else:

      # get rid of all unnecessary characters
      tmp=re.findall(r'([a-z\d]+)',conf)
      conf=''
      for t in tmp: conf+=t

      # make temporary configuration structure
      ss=re.findall(r'([a-z])',conf)
      sslist=[]
      for s in ss: sslist.append(['',s,''])
      nss=len(sslist)

      # separate digits by choosing largest allowed occupation
      digits=re.findall(r'(\d+)',conf)
      sslist[0][0]=int(digits[0])     # first n value is easy
      for i in range(nss-1):
        s=sslist[i][1]
        l=SPECLETTERS.index(s.upper())
        wmax=4*l+2
        j=len(digits[i+1])-1
        while j > 0 and int(digits[i+1][:j]) > wmax: j-=1
        w=digits[i+1][:j]
        if w == '': w=1
        sslist[i][2]=int(w)
        sslist[i+1][0]=int(digits[i+1][j:])
      sslist[-1][2]=1
      if len(digits) == nss+1: sslist[-1][2]=int(digits[-1])

    # assemble sub-shell dictionary
    ssdict={}
    wmax={}
    for ss in sslist: 
      s=str(ss[0])+ss[1]
      w=int(ss[2])
      if s not in ssdict: ssdict[s]=0
      ssdict[s]+=w
      wmax[s]=4*SPECLETTERS.index(ss[1].upper())+2
      if ssdict[s] > wmax[s]: 
        msg='maximum occupancy exceeded; conf: %s'%conf
        raise ErrorRAP(msg)
    sss=ssdict.keys()
    sss.sort(rlf_sssorter)
    first=rlf_ssindex(sss[0])
    if idxhole and idxhole >= first:
      if sshole not in ssdict: ssdict[sshole]=0
      ssdict[sshole]+=whole
      if sshole in wmax and ssdict[sshole] > wmax[sshole]: 
        msg='maximum occupancy exceeded; conf: %s'%conf
        raise ErrorRAP(msg)
      sss=ssdict.keys()
      sss.sort(rlf_sssorter)
    tconf=[]
    for s in sss: tconf.append((s,ssdict[s]))

    # deal with hole-notation
    idxfillin=[]
    if idxhole and idxhole < first:
      ssdict[sshole]=whole
      for i in range(idxhole+1,first):
        ss=rlf_subshell(i)
        ssdict[ss]=rlf_ssoccupancy(ss)
        idxfillin.append(rlf_ssindex(ss))
    sss=ssdict.keys()
    sss.sort(rlf_sssorter)

    # fill in full sub-shells
    first=rlf_ssindex(sss[0])
    for i in range(1,first):
      ss=rlf_subshell(i)
      ssdict[ss]=rlf_ssoccupancy(ss)
      idxfillin.append(rlf_ssindex(ss))

    # prepare output
    out={}
    for ss in ssdict: 
      if ssdict[ss] > 0: out[rlf_ssindex(ss)]=ssdict[ss]

    # check if configuration has the right number of electrons
    if ne is None: return out
    nediff=sum(ssdict.values())-ne
    if nediff == 0: return out
    if nediff < 0:
      msg='specified number of electrons exceeded; conf: %s'%conf
      raise ErrorRAP(msg)
    if nediff%2 == 1: 
      msg='cannot satisfy given number of electrons; conf: %s'%conf
      raise ErrorRAP(msg)

    # look for single sub-shell outside of a core
    tmp=[ssdict[ss] for ss in sss]
    if tmp.count(0) == len(tmp)-1:
      for ss in sss:
        if ssdict[ss] > 0: break
      nechk=ne-ssdict[ss]
      tout={}
      tout[rlf_ssindex(ss)]=ssdict[ss]
      if nechk == 0: return tout
      i=1
      tot=0
      okay=False
      while 1:
        w=rlf_ssoccupancy(i)
        tot+=w
        tout[i]=w
        if tot == nechk: okay=True
        if tot >= nechk: break
        i+=1
      if okay: return tout

    # search for unspecified empty sub-shells using a combinatorial
    # approach.  select matches based on fewest sub-shells and
    # highest set of sub-shells (2s > 1s and 2s 2p > 1s 2p).
    idxfillin.sort()
    idxfillin=tuple(idxfillin)
    wlist={}
    for idx in idxfillin:
      wlist[idx]=rlf_ssoccupancy(idx)
    match=False
    for n in range(1,len(idxfillin)+1):
      for comb in iter_combinations(idxfillin,n):
        wtot=0
        for idx in comb: wtot+=wlist[idx]
        if wtot == nediff: 
          rank=sum(comb)
          if match is False or rank > sum(match): match=comb
      if match is not False: break
    if match is False: 
      msg='cannot satisfy given number of electrons; conf: %s'%conf
      raise ErrorRAP(msg)
    idxmin=min(match)
    idxmax=max(match)
    for i in range(idxmin,idxmax+1):
      if i in match: del out[i]
    return out

# ------------------------------------------------------------------------------

  def copy(self):
    """ Return a copy. """
    return copy.deepcopy(self)

# ------------------------------------------------------------------------------

  def label(self,showclosed=False):
    if self.num_electrons() == 0: return ''
    out=[]
    if showclosed:
      for ss in self.subshells(): out.append(ss+str(self.occupancy(ss)))
    else:
      imax=max(self.v_conf.keys())
      icore=1
      while icore in self.v_conf and \
            self.v_conf[icore] == rlf_ssoccupancy(icore):
        icore+=1
      if icore > imax: 
        icore=imax
      elif icore not in self.v_conf:
        out.append(rlf_subshell(icore)+'0')
      for i in self.subshells(indices=True):
        w=self.occupancy(i)
        if i < icore: continue
        out.append(rlf_subshell(i)+str(w))
    out=' '.join(out)
    return out

# ------------------------------------------------------------------------------

  def num_electrons(self):
    """ Return total number of electrons. """
    return sum(self.v_conf.values())

# ------------------------------------------------------------------------------

  def subshells(self,indices=False):
    """ Return sorted list of occupied sub-shells. """
    ss=self.v_conf.keys()
    ss.sort()
    if indices: return ss
    return map(rlf_subshell,ss)

# ------------------------------------------------------------------------------

  def occupancy(self,ss):
    """ 
    Return occupancy of sub-shell, ss (which can be index or string form. 
    """
    sss=ss
    if not isinstance(sss,int): sss=rlf_ssindex(sss)
    if sss not in self.v_conf: return 0
    return self.v_conf[sss]

# ------------------------------------------------------------------------------

  def parity(self):
    """ Return parity of configuration. """
    p=1
    for i in self.subshells():
      (n,l)=rlf_ssdecode(i)
      if l%2 == 0: continue
      p*=(-1)**self.occupancy(i)
    return (1-p)/2

# ------------------------------------------------------------------------------

  def weight(self):
    """ Return statistical weight of configuration. """
    wgt=1
    for ss in self.subshells():
      w=self.occupancy(ss)
      (n,l)=rlf_ssdecode(ss)
      wgt*=rlf_binom(4*l+2,w)
    return int(wgt)

# ------------------------------------------------------------------------------

  def valence_subshell(self):
    """ Return valence sub-shell. """
    tss=max(self.subshells(indices=True))
    return rlf_subshell(tss)

# ------------------------------------------------------------------------------

  def pop_electron(self,ss=None):
    """ 
    Return a copy of the configuration where an electron from the given 
    sub-shell has been removed (last sub-shell by default).
    """
    out=copy.deepcopy(self)
    tss=ss
    if tss is None: 
      tss=max(out.subshells(indices=True))
    else:
      if not isinstance(tss,int): tss=rlf_ssindex(tss)
    if tss not in out.subshells(indices=True): 
      msg='specified sub-shell not present: %s'%ss
      raise ErrorRAP(msg)
    out.v_conf[tss]-=1
    if out.v_conf[tss] == 0: del out.v_conf[tss]
    return out

# ------------------------------------------------------------------------------

  def add_electron(self,ss=None):
    """ 
    Return a copy of the configuration where an electron has been added to the
    given sub-shell (valence sub-shell by default).
    """
    out=copy.deepcopy(self)
    tss=ss
    if tss is None: 
      tss=max(out.subshells(indices=True))
    else:
      if not isinstance(tss,int): tss=rlf_ssindex(tss)
    if tss not in out.subshells(indices=True):
      out.v_conf[tss]=1
    else:
      if out.v_conf[tss] < rlf_ssoccupancy(tss):
        out.v_conf[tss]+=1
      else:
        if ss is None:
          out.v_conf[tss+1]=1
        else:
          msg='specified sub-shell is full: %s'%ss
          raise ErrorRAP(msg)
    return out

# ------------------------------------------------------------------------------

  def difference(self,conf):
    """ 
    Return dictionary of sub-shells and difference of occupation numbers.
    The sign of the difference is according to self - conf.
    """
    sss=self.subshells(indices=True)
    osss=conf.subshells(indices=True)
    out={}
    for ss in sss:
      w=self.occupancy(ss)
      if ss not in osss:
        out[rlf_subshell(ss)]=w
      else:
        dw=w-conf.occupancy(ss)
        if dw != 0: out[rlf_subshell(ss)]=dw
    for ss in osss:
      if ss in sss: continue
      out[rlf_subshell(ss)]=-conf.occupancy(ss)
    return out

# ------------------------------------------------------------------------------

  def to_jconfig(self):
    """ Return list of possible J configurations. """
    clist=[[]]
    for i in self:
      (n,l)=rlf_ssdecode(i)
      w=self[i]

      if l == 0:
        jss1=rlf_jssencode(n,l,+1)
        t1=jss1+str(w)
        tout=[]
        for tconf in clist:
          tw=tconf[:]+[t1]
          tout.append(tw)
        clist=tout[:]
        continue

      jss1=rlf_jssencode(n,l,-1)
      jss2=rlf_jssencode(n,l,+1)
      s1max=min(w,2*l)
      s1min=max(0,w-2*l-2)

      tmp=[]
      for is1 in range(s1min,s1max+1): 
        t1=jss1+str(is1)
        t2=jss2+str(w-is1)
        tmp.append([t1,t2])

      tout=[]
      for tconf in clist:
        for tt in tmp: 
          tw=tconf[:]+tt
          tout.append(tw)
      clist=tout[:]

    out=[]
    for cc in clist:
      cstr=' '.join(cc)
      out.append(rap_jconfig(cstr))
    return out

# ------------------------------------------------------------------------------

  def allowed_terms(self,encoded=True):
    """ 
    Return list of allowed LS terms for configuration.  Return value is a
    dictionary where the keys are the term strings and the values hold the
    number of terms having the same designation.  Set encoded=False to have
    each keys be a tuple containing 2S+1, L, and parity.
    """
    termlist={}
    termlist[(1,0,0)]=1
    for i in self.v_conf:
      w=self.v_conf[i]
      (n,l)=rlf_ssdecode(i)
      tl=rlf_equiv_lsterms(l,w)
      termlist=rlf_lsterm_couple(termlist,tl)
    if not encoded: return termlist

    out={}
    for (g,l,p),n in termlist.iteritems():
      t=str(g)+SPECLETTERS[l].upper()
      out[t]=n
    return out

# ------------------------------------------------------------------------------

  def eissner(self):
    """ Return configuration label in Eissner notation. """
    out=''
    for ss in self.subshells(indices=True):
      x=rlf_ss2eissner(ss)
      q=50+self[ss]
      out+=str(q)+x
    q=int(out[0])-5
    if q == 0:
      out=out[1:]
    else:
      out[0]=str(q)
    return out

